package ch.epfl.lara.pm.analyze;

import scala.collection.mutable.HashMap

import ch.epfl.lara.pm.ast._
import ch.epfl.lara.pm.util.MyImplicitConversions
import ch.epfl.lara.pm.global.ClassTable

trait InferSetConstraints extends MyImplicitConversions {

  
  type Constraint = (String,Formula)
  
  
  private def filter(set: String)(compare: String) = set == compare 
  
  /**
   * From the class declaration infer all possible constraints on the hierarchy.
   * First, find all top class definition. This will determine the number of hierarchies.
   * Next, split classes according to the hierarchy they belong to.
   * Finally, infer all constraints from the hierarchy.
   * Assumpion: class and fields declaration are visibles. This is generally true
   *            for languages like Java and Scala.
   *
   * @param filter : Filter only classes that belong to a given hierarchy
   * @param klasses: List of class declaration
   * @return 	   : The list of constraint. Each constraint is represented as a pair
   *		     whose first value is the root class of the hierarchy to which the constraint 
   *		     (second value) apply.
   */
  protected def buildKlassConstraints(filter: String => Boolean)(klasses: List[ClassDef]): List[Constraint] = {
    val roots = for(klass <- klasses if filter(klass.name) && klass.isRoot) yield klass
    val subclasses = klasses.diff(roots)
    val hierarchies = for(root <- roots) yield {
      (root,
          for(klass <- subclasses if klass.isSubtype(root)) yield klass
      )
    }
    
    for((root,subclasses) <- hierarchies) yield {
      val partial = And(
          //1. ALL x. x  <: root -> x <: super(x)
          subclasses.map(klass => SubsetEq(klass.name,klass.superc)).reduceRight[Formula]((set,cnf) => And(set,cnf)),
              //2. Intersection over all subtypes_i <: empty
              SubsetEq(subclasses.map(klass => Set(klass.name)).reduceRight[Term]((set,in) => Intersection(set,in)),EmptySet)
          
      )
      //3. By definition: the union of the immediate subclasses of an abstract class are a superset(:>) of the set defined by the abstract class
      //   This information is necessary to verify completeness in presence of abstract classes
      val (abstrclss,clss) = (root::subclasses).partition(klass => klass.isAbstract)
      val covering = for(abstr<-abstrclss) yield {
        SubsetEq(abstr.name,klasses.filter(klass => klass.superc == abstr.name).map(klass => Set(klass.name)).reduceRight[Term]((a,b) => Union(a,b)))
      }
      val formula = covering.foldLeft[Formula](partial)((acc,form) => And(acc,form))
      (root.name,formula) 
    }
   
  }
  
  /**
   * From the objects declaration infer all possible constraints on the extractors.
   * First, get the type of the extractor's argument and use the ClassTable to get the class declaration.
   * Next, use the class declaration to obtain the root class of the hierarchy.
   * Finally, create the constraint that is simply "ObjectName <: ArgExtractorType" and update the map.
   *
   * @param objects: List of object declaration. Each object contains only an extractor declaration.
   * @return 	   : The list of constraint. Each constraint is represented as a pair
   *		     whose first value is the root class of the hierarchy to which the constraint 
   *		     (second value) apply.
   */
  protected def buildObjectConstraints(filter: String => Boolean)(objects: List[ObjectDef]): List[Constraint] = {
    val map = new HashMap[String,Formula]
    for(ObjectDef(name,extr)<-objects)  {
      val kdef = ClassTable.lookupClass(extr.getParamtype)
      val ktop = kdef.getRootName
      val constr = SubsetEq(name,Set(kdef.name))
      if(filter(ktop))
        map.get(ktop) match {
          case None => map += ((ktop,constr))
          case Some(formula) => map.update(ktop,And(constr,formula)) 
        }
    }
    map.elements.toList
  }
  
  
  /**
   * Merge two list of constraints.
   *
   * @param constr1: The first list of constraint
   * @param constr2: The second list of constraint
   * @return       : The merged list of constraint 
   */
  protected def merge(constr1: List[Constraint], constr2: List[Constraint]): List[Constraint] =  {
      for((root,c1) <- constr1) yield {
        constr2.find(pair => pair._1 == root) match {
          case None => (root,c1)
          case Some((_,c2)) => (root,And(c1,c2))
        }
      }
  }
  
  
  /**
   * Take the IR of source (the IR starts with a Program node) and the hierarchy whose constraint 
   * we need to verify some PM expression. It return the generate formula, that is nothing but
   * the conjunction of the inferred constraints on classes, objects(extractos) and the ones that the user
   * gave in the specification
   *
   * @param set    : It is the top-level class of the hierarchy. It acts as filter.
   * @param program: The top node of the IR
   * @return       : The inferred Formula. 
   */
  def constraintsOf(set: String, program: Tree): Option[Formula] = program match {
    case Program(constr,objs,clss,_) =>
      constr.filter(pair => pair._1 == set) match {
        case Nil if set != ClassTable.root =>
          val klass = ClassTable.lookupClass(set)
          if(klass.isRoot)
            None
          else
            constraintsOf(klass.superc,program)
            
        case Nil if set == ClassTable.root => None
            
        case constr => 
          val infered = merge(buildObjectConstraints(filter(set))(objs),
                              buildKlassConstraints(filter(set))(clss)
                             ).filter(pair => pair._1 == set).map(pair => pair._2)  
          val userconstr = constr.map(pair => pair._2).foldLeft[Formula](BTrue)((acc,formula) => And(acc,formula))
          Some(infered.foldLeft[Formula](userconstr)((acc,form) => And(acc,form)))
      }
      
    case _ => error("IR top-level node must be a Program")
  }
  
}
